import numpy as np
import torch
from scipy import io
from torch.utils.data import Dataset

from common_utils.utils import l2_normalize

device = 'cuda' if torch.cuda.is_available() else 'cpu'


class Dataset(Dataset):
    '''
    Inputs
        data_dir String 数据集根目录
        dataset String 数据集
        phase String train, val, test
    '''

    def __init__(self, data_dir, dataset, phase,is_gzsl = True):
        self.data_dir = data_dir
        self.dataset = dataset
        self.phase = phase

        res101 = io.loadmat(data_dir +dataset+'/'+ 'res101.mat')
        att_splits = io.loadmat(data_dir  +dataset+'/'+ 'att_splits.mat')
        train_loc = 'train_loc' # czsl 训练集位置
        val_loc = 'val_loc' # czsl 测试集位置
        trainval_loc = 'trainval_loc' # gzsl 训练集 位置
        test_seen_loc = 'test_seen_loc' # gzsl 可见类 位置
        test_unseen_loc = 'test_unseen_loc' # gzsl 不可见类位置
        # 取出训练好的特征
        feat = torch.tensor(res101['features'].T,dtype=torch.float32)
        feat = l2_normalize(feat, dim=1) # l2正则化
        self.feat =feat
        self.feat_dim = feat.shape[1] # 特征维度
        # czsl 训练集图片特征，由于train_loc是从1开始存储的，所以这里取出来时，下标-1
        self.train_data = feat[np.squeeze(att_splits[train_loc] - 1).tolist()]
        # czsl 测试集图片特征
        self.val_data = feat[np.squeeze(att_splits[val_loc] - 1).tolist()]
        # gzsl 训练集图片特征
        self.X_trainval_gzsl = feat[np.squeeze(att_splits[trainval_loc] - 1).tolist()]
        # gzsl 测试集可见类图片特征
        self.test_seen_data = feat[np.squeeze(att_splits[test_seen_loc] - 1).tolist()]
        # gzsl 测试集不可见类图片特征
        self.test_unseen_data = feat[np.squeeze(att_splits[test_unseen_loc] - 1).tolist()]
        print('Tr:{}; Val:{}; Ts:{}\n'.format(self.train_data.shape[0], self.val_data.shape[0], self.test_unseen_data.shape[0]))

        # 标签
        labels = torch.tensor(res101['labels'], dtype=torch.long)
        # 全部数据集便签，由于便签是从1开始存储的，所以这里取出来时，下标-1，方便程序使用
        self.labels_all = torch.squeeze(labels)-1
        # 获取图片位置对应的标签
        self.labels_train = torch.squeeze(labels[np.squeeze(att_splits[train_loc] - 1).tolist()]) -1
        self.labels_val = torch.squeeze(labels[np.squeeze(att_splits[val_loc] - 1).tolist()])-1
        self.labels_trainval_gzsl = torch.squeeze(labels[np.squeeze(att_splits[trainval_loc] - 1).tolist()])-1
        self.labels_test_seen = torch.squeeze(labels[np.squeeze(att_splits[test_seen_loc] - 1).tolist()])-1
        self.labels_test_unseen = torch.squeeze(labels[np.squeeze(att_splits[test_unseen_loc] - 1).tolist()])-1
        self.labels_test = torch.cat((self.labels_test_seen,self.labels_test_unseen), dim=0)

        # unique函数去除其中重复的元素，并按元素由大到小返回一个新的无元素重复的元组或者列表
        # 真实标签 [ 1  2  3  6  8 10 11 13 14 17 18 19 22 27 28 32 33 35 36 37 38 42 43 44, 45 46 49]
        self.train_classes_seen = np.unique(self.labels_train)
        # [ 4  5 12 15 16 20 21 25 26 29 39 40 48]
        self.val_classes_seen = np.unique(self.labels_val)
        # [ 1  2  3  4  5  6  8 10 11 12 13 14 15 16 17 18 19 20 21 22 25 26 27 28, 29 32 33 35 36 37 38 39 40 42 43 44 45 46 48 49]
        self.trainval_classes_seen = np.unique(self.labels_trainval_gzsl)
        # [ 1  2  3  4  5  6  8 10 11 12 13 14 15 16 17 18 19 20 21 22 25 26 27 28, 29 32 33 35 36 37 38 39 40 42 43 44 45 46 48 49]
        self.test_classes_seen = np.unique(self.labels_test_seen)
        # [ 7  9 23 24 30 31 34 41 47 50]
        self.test_classes_unseen = np.unique(self.labels_test_unseen)
        # [ 1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24, 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48, 49 50]
        self.test_classes = np.unique(self.labels_test) # All Classes of the dataset

        i = 0
        for labels in self.train_classes_seen:
            self.labels_train[self.labels_train == labels] = i
            i += 1

        j = 0
        for labels in self.val_classes_seen:
            self.labels_val[self.labels_val == labels] = j
            j += 1

        m = 0
        for labels in self.trainval_classes_seen:
            self.labels_trainval_gzsl[self.labels_trainval_gzsl == labels] = m
            m += 1
        # k = 0
        # for labels in self.test_classes_seen:
        #     self.labels_test_seen[self.labels_test_seen == labels] = k
        #     k += 1
        # l = 0
        # for labels in self.test_classes_unseen:
        #     self.labels_test_unseen[self.labels_test_unseen == labels] = l
        #     l += 1
        # 获取已经归一化的属性矩阵，50*85, 50个分类，每个分类85个属性
        self.sig =torch.tensor( att_splits['att'].T, dtype=torch.float32)
        self.attr_dim = self.sig.shape[1]
        # Shape -> (Number of attributes, Number of Classes)
        self.train_attrs = self.sig[(self.train_classes_seen ).tolist()]
        self.val_attrs = self.sig[(self.val_classes_seen ).tolist()]
        self.trainval_attrs = self.sig[(self.trainval_classes_seen ).tolist()]
        self.test_attrs_seen = self.sig[(self.test_classes_seen ).tolist()]
        self.test_attrs_unseen = self.sig[(self.test_classes_unseen ).tolist()]
        self.test_attrs = self.sig[(self.test_classes).tolist()]

        if self.phase == 'train':
            self.data = self.train_data
            self.labels = self.labels_train
            self.attrs = self.train_attrs
        elif self.phase == 'val':
            self.data = self.val_data
            self.labels = self.labels_val
            self.attrs = self.val_attrs
        elif  self.phase == 'test':
            self.data = self.test_unseen_data
            self.labels = self.labels_test_unseen
            self.attrs = self.test_attrs_unseen
        elif self.phase == 'trainval':
            self.data = self.X_trainval_gzsl
            self.labels = self.labels_trainval_gzsl
            self.attrs = self.trainval_attrs
        elif self.phase == 'gzsql_test':
            self.data = self.test_unseen_data+self.test_seen_data
            self.labels = self.labels_test_unseen+self.labels_test_seen
            self.attrs = self.test_attrs_unseen+self.test_attrs_seen
        elif  self.phase == 'all':
            self.data = self.train_data + self.val_data + self.test_unseen_data
            self.labels = labels
            self.attrs = self.sig
        else:
            raise NotImplemented
        ss = "看下数据集"
    def __getitem__(self, index):
        img = self.data[index]
        label = self.labels[index]
        attr_id = self.attrs[label]
        return img, label, attr_id

    def __len__(self):
        return len(self.data)